#!/usr/bin/env python3
# Author: Armit
# Create Time: 2023/01/07 

# Masked Language Model for chain sequence

from re import compile as Regex
import tkinter as tk
import tkinter.ttk as ttk
import tkinter.font as tkft
import tkinter.messagebox as tkmsg
from argparse import ArgumentParser
from threading import Thread
from collections import defaultdict, Counter
from typing import List
from traceback import print_exc, format_exc

import matplotlib.pyplot as plt
from matplotlib import ticker
from matplotlib.offsetbox import AnchoredText
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg

from data import load_pdb_data

WINDOW_TITLE    = 'MLM for AA sequence'
WINDOW_SIZE     = (720, 550)
TEXT_WIDTH      = 120
TEXT_FONT_SIZE  = 24
QUERY_LIMIT     = 5
QUERY_MARK      = '_'
QUERY_EXAMPLE   = 'A_C_CG'
NOANS_MARK      = '?'

class App:

  def __init__(self):
    self.chains: List[str] = None

    self.init()
    self.setup_gui()

    try:
      self.wnd.mainloop()
    except KeyboardInterrupt:
      self.wnd.destroy()
    except: print_exc()

  def init(self):
    residue_names, sequences, _ = load_pdb_data()
    self.chains = list(set(sequences))
    self.class_labels = residue_names
  
  def setup_gui(self):
    # window
    wnd = tk.Tk()
    W, H = wnd.winfo_screenwidth(), wnd.winfo_screenheight()
    self.loc_modal = f'+{W//2}+{H//2}'
    w, h = WINDOW_SIZE
    wnd.geometry(f'{w}x{h}+{(W-w)//2}+{(H-h)//2}')
    wnd.resizable(False, False)
    wnd.title(WINDOW_TITLE)
    self.wnd = wnd

    # modal
    self.modal = tk.Toplevel(wnd)
    ttk.Label(self.modal, text='Matching database, please wait...', foreground='red').pack()
    self._modal('hide')

    # top: query
    frm1 = ttk.LabelFrame(wnd, text='Query')
    frm1.pack(side=tk.TOP, anchor=tk.N, expand=tk.YES, fill=tk.X)
    if True:
      frm11 = ttk.Label(frm1)
      frm11.pack(side=tk.TOP)
      if True:
        self.var_query = tk.StringVar(frm11, value=QUERY_EXAMPLE)
        ent = ttk.Entry(frm11, textvariable=self.var_query, width=TEXT_WIDTH, justify=tk.CENTER, 
                        font=tkft.Font(frm11, family='Lucida Grande', size=TEXT_FONT_SIZE))
        ent.bind('<Return>', lambda evt: self.predict())
        ent.pack(expand=tk.YES, fill=tk.X)

    # bottom: plot
    frm2 = ttk.LabelFrame(wnd, text='Freq')
    frm2.pack(expand=tk.YES, fill=tk.BOTH)
    if True:
      frm21 = ttk.Frame(frm2)
      frm21.pack(expand=tk.YES, fill=tk.BOTH)
      if True:
        fig, ax = plt.subplots()
        fig.tight_layout()
        cvs = FigureCanvasTkAgg(fig, frm21)
        cvs.get_tk_widget().pack(expand=tk.YES, fill=tk.BOTH)
        self.fig, self.ax, self.cvs = fig, ax, cvs
        self._plot_show(init=True)

  def _modal(self, what):
    if what == 'show':
      self.modal.geometry(self.loc_modal)
      self.modal.update()
      self.modal.deiconify()
      self.modal.grab_set_global()
    if what == 'hide':
      self.modal.grab_release()
      self.modal.withdraw()

  def _plot_show(self, init=False):
    self.ax.xaxis.set_major_locator(ticker.FixedLocator(range(len(self.class_labels))))
    self.ax.set_xticklabels(self.class_labels)
    self.ax.set_xlim(-1, len(self.class_labels))
    if not init: self.fig.legend()
    self.cvs.draw()

  def predict(self):
    query_string = self.var_query.get().strip().upper()
    print(f'>> query: {query_string}')

    if not query_string: return
    if len(query_string) <= 3: tkmsg.showerror('Error', 'Query string too short!') ; return
    
    n_query = query_string.count(QUERY_MARK)
    if n_query <= 0: tkmsg.showerror('Error', f'No query marks {QUERY_MARK!r} found!') ; return
    if n_query > QUERY_LIMIT: tkmsg.showerror('Error', f'Two many query marks, the upper limit is {QUERY_LIMIT}!') ; return
    
    overflow = set(query_string) - set(self.class_labels) - { QUERY_MARK }
    if len(overflow) > 0:
      tkmsg.showerror('Error', f'Unknown token in query string: {overflow}!') ; return

    try:
      QUERY_REGEX = Regex(query_string.replace(QUERY_MARK, '(\w)'))
    except:
      print_exc()
      tkmsg.showerror('Error', 'Ill-formated query string, cannot compile pattern!')
      return

    def task():
      self._modal('show')
      try:
        joint_matches = []    # 联合分布，记录所有共现的query解
        matches = [defaultdict(int) for _ in range(n_query)]    # matches[query][residue] = support
        for chain in self.chains:
          m = QUERY_REGEX.findall(chain)
          if not m: continue
          joint_matches.extend(m)
          for r in m:
            if n_query == 1: r = (r,)
            for i in range(n_query):
              matches[i][r[i]] += 1
        freqs = [[preds[lbl] for lbl in self.class_labels] for preds in matches]   # freqs[query][support, ..]

        # 给出联合分布频数最大的解
        if len(joint_matches) == 0:
          solut = None
          supp = cnt = 0
        else:
          solut, cnt = Counter(joint_matches).most_common(1)[0]
          supp = cnt / len(joint_matches)
        pred = []
        q = 0
        for r in query_string:
          if r == QUERY_MARK:
            r = solut[q] if solut else NOANS_MARK
            q += 1
          pred.append(r)
        assert q == n_query
        pred = ''.join(pred)

        pred_str = f'predict: {pred}'
        supp_str = f'support: {cnt} ({supp:.3%})'
        result = f'{pred_str}\n{supp_str}'
        print(f'>> {pred_str}, {supp_str}')

        # 作图只能给出独立假设下每个query的频数分布 (实际上当然不独立)
        self.ax.clear()
        for i, freq in enumerate(freqs):
          self.ax.plot(freq, label=i)
        self.ax.add_artist(AnchoredText(result, loc=2, alpha=0.7))
        self._plot_show()
      except:
        self._modal('hide')
        info = format_exc()
        print(info)
        tkmsg.showerror('Error', info)
      finally:
        self._modal('hide')

    Thread(target=task).start()


if __name__ == '__main__':
  parser = ArgumentParser()
  parser.add_argument('--dedup', action='store_true')
  args = parser.parse_args()
  
  App()
